/***************************************************************************
 *   Copyright (C) 2010 by programmeboy                                    *
 *   programmeboy@gmail.com                                                *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.                                        *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             * 
 ***************************************************************************/

#include <ntddk.h>

#include "cpu.h"
#include "vt.h"
#include "vmx.h"
#include "bp.h"

BREAKPOINT_INFO gBpInfo[MAX_BREAK_POINT] = {0};
BOOLEAN gbIsSingleDebug = FALSE;

VOID
BPInit()
{
	ULONG i;
	for (i = 0; i < MAX_BREAK_POINT; i++){
		RtlZeroMemory(&gBpInfo[i], sizeof(BREAKPOINT_INFO));
	}
}

BOOLEAN 
BPIsSingleDebuging(
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    return gbIsSingleDebug == TRUE;
}

VOID 
BPSetSingleDebuging(
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    gbIsSingleDebug = TRUE;
}

VOID 
BPDisSingleDebuging(
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    gbIsSingleDebug = FALSE;
}

BOOLEAN
BPIsSettingBreakPoint(
		ULONG Address
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
	ULONG i;
	for (i = 0; i < MAX_BREAK_POINT; i++){
		if (gBpInfo[i].VirtualAddress == Address && \
			gBpInfo[i].IsSet == TRUE){
			return TRUE;
		}
	}
	return FALSE;
}

PBREAKPOINT_INFO
BPGetBreakPointInfo(
	ULONG Address
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
	ULONG i;
	for (i = 0; i < MAX_BREAK_POINT; i++){
		if (gBpInfo[i].VirtualAddress == Address && \
			gBpInfo[i].IsSet == TRUE){
			return &gBpInfo[i];
		}
	}
	return NULL;
};

BOOLEAN
BpSetBreakPoint(
	ULONG Address
)
/*++

Routine Description:

    Set break point

Arguments:

	

Return Value:

--*/
{
	ULONG i;
	ULONG MyCR3, OldCR3;
	MyCR3 = _ReadVMCS(GUEST_CR3);
	
	SET_CR3(MyCR3, OldCR3);
	if (!MmIsAddressValid((PVOID)Address)){
		return FALSE;
	}
	for (i = 0; i < MAX_BREAK_POINT; i++){
		if (gBpInfo[i].IsSet == FALSE){

			gBpInfo[i].SavedByte = *(PUCHAR)Address;
			*(PUCHAR)Address = INT3_CODE;
			gBpInfo[i].VirtualAddress = Address;
			gBpInfo[i].IsSet = TRUE;
			break;
		}
	}
	SET_CR3(OldCR3, MyCR3);
	if (i == MAX_BREAK_POINT){
		return FALSE;
	}
	return TRUE;
}

VOID
BpDelBreakPoint(
	ULONG Address
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
	ULONG i;
	ULONG MyCr3, OldCR3;
	MyCr3 = _ReadVMCS(GUEST_CR3);
	
	SET_CR3(MyCr3, OldCR3);
	if (!MmIsAddressValid((PVOID)Address)){
		return;
	}
	for (i = 0; i < MAX_BREAK_POINT; i++){
		if (gBpInfo[i].IsSet == TRUE && \
			gBpInfo[i].VirtualAddress == Address){
			
			*(PUCHAR)Address = gBpInfo[i].SavedByte;
			gBpInfo[i].IsSet = FALSE;
			break;
			
		}
	}
	SET_CR3(OldCR3, MyCr3);
}

BOOLEAN
BPRestoreProcessCode(
	ULONG Address
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
	ULONG i;
	ULONG MyCr3, OldCR3;
	MyCr3 = _ReadVMCS(GUEST_CR3);

	SET_CR3(MyCr3, OldCR3);
	if (!MmIsAddressValid((PVOID)Address)){
		return FALSE;
	}
	for (i = 0; i < MAX_BREAK_POINT; i++){
		if (gBpInfo[i].IsSet == TRUE && \
			gBpInfo[i].VirtualAddress == Address){

				*(PUCHAR)Address = gBpInfo[i].SavedByte;
				break;
		}
	}
	SET_CR3(OldCR3, MyCr3);
	return TRUE;
}

BOOLEAN
BPSetProcessCode(
	ULONG Address
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
	ULONG i;
	ULONG MyCr3, OldCR3;
	MyCr3 = _ReadVMCS(GUEST_CR3);

	SET_CR3(MyCr3, OldCR3);
	if (!MmIsAddressValid((PVOID)Address)){
		return FALSE;
	}
	for (i = 0; i < MAX_BREAK_POINT; i++){
		if (gBpInfo[i].IsSet == TRUE && \
			gBpInfo[i].VirtualAddress == Address){
				*(PUCHAR)Address = INT3_CODE;
				break;
		}
	}
	SET_CR3(OldCR3, MyCr3);
	return TRUE;
}

NTSTATUS
BPSetSingleStepDebug(
    PCPU_VM_CONTEXT pCPU
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    ULONG GuestRFlag;
    ULONG InsLen;
    ULONG CurrentEIP;
    PSINGLE_STEP_INFO pSingled;

    if (!BPIsSingleDebuging()){
        BPSetSingleDebuging();
        
        //
        //  Clear the single step info
        //
        
        if (pCPU->SingleStepInfo != NULL){
            ExFreePool(pCPU->SingleStepInfo);
        }
        
        //
        //  Reset the single step info
        //
        
        pCPU->SingleStepInfo = ExAllocatePool(NonPagedPool, 
                    sizeof(SINGLE_STEP_INFO));
        RtlZeroMemory(pCPU->SingleStepInfo, sizeof(SINGLE_STEP_INFO));
        
    }
    
    pSingled = pCPU->SingleStepInfo;

    //
    //  update context
    //
    
    pSingled->LastEip = _ReadVMCS(GUEST_RIP);
    pSingled->RFlag = _ReadVMCS(GUEST_RFLAGS);
    pSingled->LastInsLen = _ReadVMCS(VM_EXIT_INSTRUCTION_LEN);

    //
    //  Set the TF of RFlag
    //

    GuestRFlag = _ReadVMCS(GUEST_RFLAGS);
    if (!(GuestRFlag & FLAGS_TF_MASK)){
        GuestRFlag |= FLAGS_TF_MASK;
        _WriteVMCS(GUEST_RFLAGS, GuestRFlag);
    }
    
    return STATUS_SUCCESS;
      
}

BOOLEAN
BPIsOurSingleStep(
    PCPU_VM_CONTEXT pCPU
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    ULONG CurrentEIP;
    PSINGLE_STEP_INFO pSingled;

    if (pCPU == NULL){
        return FALSE;
    }
    if (pCPU->SingleStepInfo == NULL){
        return FALSE;
    }

    pSingled = pCPU->SingleStepInfo;

    //
    //  Is our's ?
    //
    
    CurrentEIP = _ReadVMCS(GUEST_RIP);
    if ((pSingled->LastEip + pSingled->LastInsLen) == \
        CurrentEIP){
        return TRUE;
    }
    return FALSE;
}